from keras.models import Model
from keras.layers import Input, Conv3D, BatchNormalization, Activation, MaxPooling3D, SpatialDropout3D, merge, Flatten, Dense, Dropout, Lambda
from keras.regularizers import l2
import tensorflow as tf


def vggblock(tensor, depth, nb_filters, init, regularizer):
    for _ in range(depth):
        tensor = Conv3D(nb_filters, 3, 3, 3, init=init, border_mode='same', W_regularizer=regularizer)(tensor)
        tensor = BatchNormalization()(tensor)
        tensor = Activation('relu')(tensor)
    return tensor


def vgg13_shortcuts(input_shape, base_nb_filters, weight_decay, init='he_normal', drop_rate=0.5):
    regularizer = l2(weight_decay)
    inputs = Input(shape=input_shape)
    outputs = []
    tensor = inputs
    for depth, factor in zip([2, 2, 3, 3], [1, 2, 4, 8]):
        tensor = vggblock(tensor, depth, factor * base_nb_filters, init, regularizer)
        tensor = MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2))(tensor)
        outputs.append(tensor)
        tensor = SpatialDropout3D(drop_rate / 2)(tensor)
        print(tensor.get_shape())

    nb_outputs = len(outputs)
    for i in range(nb_outputs):
        ps = 2 ** (nb_outputs - i - 1)
        outputs[i] = MaxPooling3D(pool_size=(ps, ps, ps), strides=(ps, ps, ps))(outputs[i])
    outputs = merge(outputs, mode='concat', concat_axis=-1)
    print(outputs.get_shape())

    outputs = Flatten()(outputs)
    for _ in range(2):
        outputs = Dropout(drop_rate)(outputs)
        outputs = Dense(16 * base_nb_filters, init=init, W_regularizer=regularizer)(outputs)
        outputs = BatchNormalization()(outputs)
        outputs = Activation('relu')(outputs)
    outputs = Dropout(drop_rate)(outputs)
    outputs = Dense(2, init=init, W_regularizer=regularizer)(outputs)
    outputs = Activation('softmax')(outputs)

    return Model(input=inputs, output=outputs)


def depad(target_shape):
    def helper(x):
        origin_shape = x.get_shape().as_list()[1:-1]
        starts = [0, ] + [(os - ts) // 2 for os, ts in zip(origin_shape, target_shape)] + [0, ]
        size = [-1, ] + target_shape + [-1, ]
        x = tf.slice(x, starts, size)
        return x
    return helper


def vgg13_shortcuts_v2(input_shape, base_nb_filters, weight_decay, init='he_normal', drop_rate=0.5):
    regularizer = l2(weight_decay)
    inputs = Input(shape=input_shape)
    outputs = []
    tensor = inputs
    for depth, factor in zip([2, 2, 3, 3], [1, 2, 4, 8]):
        tensor = vggblock(tensor, depth, factor * base_nb_filters, init, regularizer)
        tensor = MaxPooling3D(pool_size=(2, 2, 2), strides=(2, 2, 2))(tensor)
        outputs.append(tensor)
        tensor = SpatialDropout3D(drop_rate / 2)(tensor)
        print(tensor.get_shape())

    nb_outputs = len(outputs)
    target_shape = outputs[-1].get_shape().as_list()[1:-1]
    for i in range(nb_outputs):
        output = outputs[i]
        output = Lambda(depad(target_shape), output_shape=target_shape + output.get_shape().as_list()[-1:])(output)
        print(output.get_shape())
        output = Flatten()(output)
        print(output.get_shape())
        for _ in range(2):
            output = Dropout(drop_rate)(output)
            output = Dense(16 * base_nb_filters, init=init, W_regularizer=regularizer)(output)
            output = BatchNormalization()(output)
            output = Activation('relu')(output)
        outputs[i] = output
    outputs = merge(outputs, mode='concat', concat_axis=-1)
    outputs = Dropout(drop_rate)(outputs)
    outputs = Dense(2, init=init, W_regularizer=regularizer)(outputs)
    outputs = Activation('softmax')(outputs)

    return Model(input=inputs, output=outputs)
